from __future__ import division
import jieba
import sys
import pickle
import yaml
import operator
import os


THRESHOLD = 0.3


def pickle_load(f):
    with open(f, 'rb') as fd:
        ret = pickle.load(fd)
    return ret


class MySentences(object):
    def __init__(self, dir_name):
        self.dir_name = dir_name

    def __iter__(self):
        for f_name in os.listdir(self.dir_name):
            if f_name[0] == '.':
                continue
            print(os.path.join(self.dir_name, f_name))
            for line in open(os.path.join(self.dir_name, f_name)):
                yield line


class Classifier:
    def __init__(self, cv_file, tfidf_file, lg_file):
        self.cv, self.tfidf, self.lg = self.load_model(cv_file, tfidf_file, lg_file)

    @staticmethod
    def load_model(cv_file, tfidf_file, lg_file):
        cv = pickle_load(cv_file)
        tfidf = pickle_load(tfidf_file)
        lg = pickle_load(lg_file)

        return cv, tfidf, lg

    def classify(self, text_list):
        words = ' '.join(text_list)
        d = self.cv.transform([words])
        d = self.tfidf.transform(d)
        res = zip(self.lg.classes_, self.lg.predict_proba(d)[0])
        out = sorted(res, key=operator.itemgetter(1), reverse=True)
        return out


def get_class2_id(res, c_id_dict, c2_dict, p_id, threshold=THRESHOLD):
    out = []
    for k, p in res:
        c2_id = c2_dict[k]
        if c_id_dict[c2_id] == p_id:
            if p >= threshold:
                out.append(c2_id)
            if len(out) == 0:
                out.append(c2_id)
                break
            if p < threshold:
                break
    return out


def get_multi_class1_id(c_out, c_dict, threshold):
    res = []
    for c, prob in c_out:
        if prob >= threshold:
            res.append(c_dict[c])
        else:
            break
    return res


if __name__ == '__main__':
    config_file = sys.argv[1]
    class_id_file = sys.argv[2]
    in_dir = sys.argv[3]
    in_file = sys.argv[4]
    out_file = sys.argv[5]

    pid_dict = {}
    for k in open(in_file):
        pid = k.split('\t')[0]
        pid_dict[pid.strip()] = 1

    class1_dict = {}
    class2_dict = {}
    class_id_dict = {}
    for l in open(class_id_file):
        d = l.strip().split('\t')
        if len(d) != 3:
            continue
        c_id, p_id, c_name = d
        if p_id == '0':
            class1_dict[c_name] = c_id
        else:
            class2_dict[c_name] = c_id

        class_id_dict[c_id] = p_id

    # print(class1_dict)
    # print(class2_dict)
    # print(class_id_dict)

    conf = yaml.load(open(config_file, 'r'))
    classifiers = conf['classifiers']

    # count_vector_file = '/Users/hardy/data/gen_models/cv.pkl'
    # tfidf_file = '/Users/hardy/data/gen_models/tfidf.pkl'
    # lg_file = '/Users/hardy/data/gen_models/lg.pkl'
    cv1 = classifiers['c1']['cv']
    tfidf1 = classifiers['c1']['tfidf']
    lg1 = classifiers['c1']['lg']

    cv2 = classifiers['c2']['cv']
    tfidf2 = classifiers['c2']['tfidf']
    lg2 = classifiers['c2']['lg']
    clf1 = Classifier(cv1, tfidf1, lg1)
    clf2 = Classifier(cv2, tfidf2, lg2)

    out_fd = open(out_file, 'w+')
    data_iter = MySentences(in_dir)
    for l in data_iter:
        t = l.strip()
        d = t.split('\t', 1)

        pid = d[0].strip()
        if pid_dict.get(pid):
            continue
        # title_content = d[1] + d[2]
        if len(d) != 2:
            print(pid + '\t' + '\t' + '\n')
            continue
        seg = jieba.lcut(d[1])

        c1_out = clf1.classify(seg)
        c2_out = clf2.classify(seg)

        c1 = c1_out[0][0]
        c2 = c2_out[0][0]
        class1_id = class1_dict[c1]
        class2_id = class2_dict[c2]

        class1_ids = get_multi_class1_id(c1_out, class1_dict, THRESHOLD)
        class2_ids = []

        if not class1_ids:
            class1_ids.append(class1_id)

        for c1_id in class1_ids:
            c2_tmp = get_class2_id(c2_out, class_id_dict, class2_dict, c1_id)
            class2_ids += c2_tmp

        class1_id = ','.join(class1_ids)
        class2_id = ','.join(class2_ids)

        out = '\t'.join([pid, class1_id, class2_id])
        out_fd.write(out + '\n')

    out_fd.close()



